from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# 加载微调好的模型和分词器
model_path = "./fine_tuned_roberta_sentiment_chinese"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)

def predict_sentiment(text):
    # 预处理输入
    inputs = tokenizer(text, 
                      return_tensors="pt", 
                      truncation=True, 
                      padding=True,
                      max_length=256)
    
    # 模型预测
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_label = torch.argmax(logits, dim=1).item()
    
    # 输出结果（假设0=负面，1=正面）
    label_map = {0: "负面", 1: "正面"}
    return label_map[predicted_label]

# 交互式测试
if __name__ == "__main__":
    print("中文情感分类测试（输入q退出）")
    while True:
        text = input("\n请输入中文文本：")
        if text.lower() == 'q':
            break
        result = predict_sentiment(text)
        print(f"预测结果: {result}")